'''
Created on Aug 9, 2009
@author: mjbommar
'''

import sqlite3

try:
    import numpy
    import scipy
    import pylab
except:
    raise Exception("numpy, scipy, and pylab must all be installed")
    
from tfaces.collection.util import *

class TimeSeriesDatabase(object):
    '''
    This object reads in the SQLite output of a Time Series Experiment parsing.
    It allows for a wide range of ways to access or manipulate data.
    '''
    
    def __init__(self, dbFile, independentVariable = None):
        '''
        Constructor that takes the database file name.
        '''
        self.db = sqlite3.Connection
        self.db = sqlite3.connect(dbFile)
        
        self.parameters = []
        self.loadParameters()
        
        self.outputs = []
        self.loadOutputs()
        
        if independentVariable:
            self.independentVariable = independentVariable
            self.independentVariableIndex = self.outputs.index(independentVariable)
        else:
            self.determineIndependentVariable()
            self.independentVariableIndex = self.outputs.index(self.independentVariable)
    
    def loadParameters(self):
        '''
        Load the list of parameters from the database file.
        We use the PRAGMA table_info command to retrieve this from the 
        parameterset table.
        '''
        results = self.db.execute('PRAGMA table_info(parameterset)').fetchall()
        self.parameters = [row[1] for row in results]
        self.parameters.remove('id')
    
    def loadOutputs(self):
        '''
        Load the list of outputs from the database file.
        We use the PRAGMA table_info command to retrieve this from the 
        runs table, ignoring the seed and parametersetid fields.
        '''
        results = self.db.execute('PRAGMA table_info(runs)').fetchall()
        self.outputs = [row[1] for row in results]
        self.outputs.remove('seed')
        self.outputs.remove('parametersetid')
        
    def determineIndependentVariable(self):
        '''
        Attempt to determine the independent variable in the 
        set of outputs.  If none is found, raise an exception.
        '''
        possibleColumns = ['time', 't', 'column0']

        for output in self.outputs:
            if output.lower() in possibleColumns:
                self.independentVariable = output
                return
        
        raise Exception("Database::determineIndependentVariable: unable to find suitable output")
    
    def describeParameters(self):
        '''
        Describe the experiment parameters.
        This function lists each parameter as well as the values it takes.
        Stars in the left column indicate that the parameter is swept.
        '''
        print 'Parameters'
        for parameter in self.parameters:
            values = self.getParameterValues(parameter)
            
            # Make sure we delimit the values properly if they are strings
            if type(toNaturalType(values[0])) == str:
                values = ['"%s"' % value for value in values]
            else:
                values = map(str, values)

            if len(values) > 1:
                print "*", "\t", parameter, "=", ",".join(values)
            else:
                print "\t", parameter, "=", ",".join(values)
        
    def describeOutputs(self):
        '''
        Describe the experiment outputs.
        This function lists the outputs.
        A star on the left indicate that the parameter is determined to be the IV.
        '''
        print 'Outputs'
        for output in self.outputs:
            if output == self.independentVariable:
                print "*", "\t", output, "<-- Independent Variable"
            else:
                print "\t", output

    def describeExperiment(self):
        '''
        Describe the parameters and outputs of the experiment.
        '''
        self.describeParameters()
        print
        self.describeOutputs() 

    def getParameterValues(self, parameter):
        '''
        Return the set of all values that a given parameter takes.
        '''
        sqlQuery = 'SELECT DISTINCT %s FROM parameterset' % parameter
        rows = [row[0] for row in self.db.execute(sqlQuery, tuple()).fetchall()]
        return rows            
    
    def getRunBySeed(self, seed, outputs = None):
        '''
        Return the output of a run selected by its seed.
        If outputs is None, then all outputs are returned for the run.
        However, outputs can be a list of the desired output fields to return.
        '''
        
        # Check to see if the list of desired output fields was passed.
        if outputs:
            for output in outputs:
                if output not in self.outputs:
                    raise Exception("Database::getRunBySeed: %s is not in the list of database outputs" % output)

            sqlQuery = 'SELECT ' + ','.join(outputs) + ' FROM runs WHERE seed = ?'
        else:
            outputs = self.outputs
            sqlQuery = 'SELECT * FROM runs WHERE seed = ?'
        
        # Sort by the IV if it is given, execute the query and return the rows    
        if self.independentVariable:
            sqlQuery += ' ORDER BY %s ASC' % self.independentVariable
            rows = self.db.execute(sqlQuery, (seed,)).fetchall()
        else:
            rows = self.db.execute(sqlQuery, (seed,)).fetchall()
        
        M = numpy.array(rows)
        return M

    def getSeedsByParameterSetID(self, id):
        '''
        Return the seeds corresponding to a given parameter set id.
        These can then be passed to getRunBySeed.
        '''
        sqlQuery = 'SELECT DISTINCT seed FROM runs WHERE parametersetid = ?'
        rows = self.db.execute(sqlQuery, (id,)).fetchall()
        
        if len(rows) > 0:
            seeds = [row[0] for row in rows]
        else:
            raise Exception("Database::getSeedsByParameterSetID: no seeds match given parameter set ID")

        return seeds
    
    def getParameterSetIDBySeed(self, seed):
        '''
        Return the parameter set ID for a given seed.
        '''
        sqlQuery = 'SELECT DISTINCT parametersetid FROM runs WHERE seed = ?'
        rows = self.db.execute(sqlQuery, (seed,)).fetchall()
        
        if len(rows) == 0:
            raise Exception("Database::getParameterSetIDBySeed: no run matching seed %s" %(seed))
        
        return rows[0][0]
        
    
    def getParameterValuesBySeed(self, seed):
        '''
        Return the actual parameter values for a given seed.
        '''
        id = self.getParameterSetIDBySeed(seed)
        sqlQuery = 'SELECT * FROM parameterset WHERE id = ?'
        rows = self.db.execute(sqlQuery, (id,)).fetchall()
        
        if len(rows) == 0:
            raise Exception("Database::getParameterValuesBySeed: no parameter set matches given seed")
        
        print dict(zip(self.parameters, map(toNaturalType, rows[0])))

        
    
    def getParameterSetByParameters(self, **constraints):
        '''
        Return the IDs of the parameter sets that match the given parameters constraints.
        Constraints can be of the form:
            namedArgument = value
            namedArgument = "<value"
            namedArgument = "<=value"
            namedArgument = ">value"
            namedArgument = ">=value"
        '''

        # Check that constraints are actually given
        if len(constraints) == 0:
            raise Exception("Database::getOutputsByConstraints: no constraints given")
        
        # Determine the valid parameter set IDs
        sqlQuery = "SELECT id FROM parameterset WHERE "

        values = []
        for (key,value) in constraints.iteritems():
            # Check to make sure that the given parameter is valid
            if key not in self.parameters:
                raise Exception("Database::getOutputsByConstraints: %s is not a valid parameter" % key)
        
            if type(value) != str:
                sqlQuery += ' %s = ? AND' % (key,)
                values.append(value)
            else:
                # Parse the possible inequalities
                if value.startswith('>='):
                    op = '>='
                    value = value[2:]
                elif value.startswith('>'):
                    op = '>'
                    value = value[1:]
                elif value.startswith('<='):
                    op = '<='
                    value = value[2:]
                elif value.startswith('<'):
                    op = '<'
                    value = value[1:]
                
                sqlQuery +=  '%s %s ? AND' % (key, op)
                values.append(value)
        
        sqlQuery = sqlQuery.strip(' AND').strip()
        rows = self.db.execute(sqlQuery, values).fetchall()
        if len(rows) == 0:
            raise Exception("Database::getOutputsByConstraints: no parameter sets match the given constraints")
        
        ids = [row[0] for row in rows]
        return ids

    def getRunsByParameterSetID(self, id, outputs = None):
        '''
        Returns a list of numpy arrays corresponding to each run.
        A list of the desired output fields can be passed to outputs, just as in getRunBySeed.
        '''
        runs = []
        
        # Get the seeds for the given parameter set
        seeds = self.getSeedsByParameterSetID(id)
        for seed in seeds:
            runs.append(self.getRunBySeed(seed, outputs))

        return runs
    
    def getAverageRunByParameterSetID(self, id, outputs = None):
        '''
        Returns an average of the runs of a parameter set.
        A list of the desired output fields can be passed to outputs, just as in getRunBySeed.
        '''
        # Get the runs for the given parameter set ID
        runs = self.getRunsByParameterSetID(id, outputs)
        
        # Create the averaged matrix
        M = numpy.zeros(runs[0].shape)
        for run in runs:
            M += run
        M /= float(len(runs))
        
        return M
    
    def getStdRunByParameterSetID(self, id, outputs = None):
        '''
        Returns the standard deviation of the values of each element of the average run array
        returned by getAverageRunByParameterSetID.
        '''
        # Get the runs for the given parameter set ID
        runs = self.getRunsByParameterSetID(id, outputs)
        M = numpy.zeros(runs[0].shape)
        
        # Create the matrix with each value 
        for i in range(M.shape[0]):
            for j in range(M.shape[1]):
                M[i,j] = numpy.std([run[i,j] for run in runs])
        
        return M
    
    def getAverageOutputBySeed(self, seed, outputs = None):
        '''
        Return the average value of an output for a given run specified by its seed.
        '''
        D = self.getRunBySeed(seed, outputs)
        return numpy.mean(D,0)
    
    def getStdOutputBySeed(self, seed, outputs = None):
        '''
        Return the standard deviation of an output for a given run specified by its seed.
        '''
        D = self.getRunBySeed(seed, outputs)
        return numpy.std(D, 0)

    def getAverageOutputByParameterSetID(self, id, outputs = None):
        '''
        Return the average value of an output for a given parameter set.
        '''
        # Get the runs for the given parameter set ID
        runs = self.getRunsByParameterSetID(id, outputs)
        
        if not outputs:
            outputs = self.outputs
        
        averages = [numpy.mean(numpy.array([run[:,i] for run in runs])) for i,output in enumerate(outputs)]
        return averages
    
    def getStdOutputByParameterSetID(self, id, outputs = None):
        '''
        Return the standard deviation of an output for a given parameter set.
        '''
        # Get the runs for the given parameter set ID
        runs = self.getRunsByParameterSetID(id, outputs)
        
        if not outputs:
            outputs = self.outputs
        
        stds = [numpy.std(numpy.array([run[:,i] for run in runs])) for i,output in enumerate(outputs)]
        return stds
    
    def getRunsByConstraints(self, outputs, **constraints):
        '''
        Return the selected outputs for all runs where the given constraints are met.
        Note that outputs must be a list of output fields or None to include all.  
        The constraints are to be given as named argument of the form:
            nA = value
            nA = "<value"
            nA = ">value"
        '''
        
        # Check to make sure that the given outputs are valid
        for output in outputs:
            if output not in self.outputs:
                raise Exception("Database::getOutputsByConstraints: %s is not a valid output" % output)
        
        # Get the parameter set IDs and then return the runs from the parameter sets
        ids = self.getParameterSetByParameters(**constraints)
        pRuns = [self.getRunsByParameterSetID(id, outputs) for id in ids]
        runs = []
        for p in pRuns:
            runs.extend(p)
        
        return runs
    
    def getAverageRunByConstraints(self, outputs, **constraints):
        '''
        Return the average of a set of runs that match the given constraints.
        See getRunsByConstraints for more details.
        
        All runs must have output of the same size.
        '''
        # Get the runs for the given parameter set ID and put them into a single list
        runs = self.getRunsByConstraints(outputs, **constraints)

        # Direct-sum the matrices and then divide
        M = numpy.zeros(runs[0].shape)
        for run in runs:
            M += run
        
        M /= float(len(runs))
        
        return M
    
    def getStdRunByConstraints(self, outputs, **constraints):
        '''
        Returns the standard deviation of the values of each element of the set of runs
        that match the given constraints.
        
        All runs must have output of the same size.
        '''
        # Get the runs for the given parameter set ID and put them into a single list
        runs = self.getRunsByConstraints(outputs, **constraints)

        # Calculate the standard deviation element-wise
        M = numpy.zeros(runs[0].shape)
        for i in range(M.shape[0]):
            for j in range(M.shape[1]):
                M[i,j] = numpy.std([run[i,j] for run in runs])
        
        return M
